# coding: utf-8


import os
import h5py
import numpy as np
from keras.preprocessing.image import ImageDataGenerator
from keras.applications import resnet50, inception_v3, xception
from keras.applications.resnet50 import ResNet50
from keras.applications.inception_v3 import InceptionV3
from keras.applications.xception import Xception
from keras.layers.core import Lambda
from keras.layers import Input, GlobalAveragePooling2D, GlobalMaxPooling2D, Dense, Dropout
from keras.models import Model
from keras.callbacks import ModelCheckpoint, EarlyStopping, ReduceLROnPlateau
from keras.optimizers import adam, Adam
from sklearn.utils import shuffle
from keras.utils import to_categorical
from keras.utils.vis_utils import plot_model
from loss_history import LossHistory
import heapq
from keras import backend as K
K.clear_session()


# -------------------------------------
# 定义参数
# -------------------------------------
EPOCHS = 1000
BATCH_SIZE = 128
LEARNING_RATE = 1e-4
EARLY_STOPPING_PATIENCE = 20
REDUCE_LR_PATIENCE = 3
CLASS_NUM = 2
DROP_RATE = 0.5


# -------------------------------------
# 文件路径
# -------------------------------------
TRAIN_DATA_PATH = './data/train/'
VALID_DATA_PATH = './data/valid/'
RESULT_PATH = './result/'
MODEL_PATH = './model/'
OUTPUT_PATH = './output/'


def mergeFinetuneModel():

    X_train = []
    X_valid = []

    filenames = [os.path.join(OUTPUT_PATH, 'inceptionv3-finetune-output.hdf5'),
                 os.path.join(OUTPUT_PATH, 'resnet50-finetune-output.hdf5'),
                 os.path.join(OUTPUT_PATH, 'xception-finetune-output.hdf5'),
                 os.path.join(OUTPUT_PATH, 'vgg16-finetune-output.hdf5')]

    for filename in filenames:
        with h5py.File(filename, 'r') as h:
            X_train.append(np.array(h['X_train']))
            X_valid.append(np.array(h['X_val']))
            y_train = np.array(h['y_train'])
            y_valid = np.array(h['y_val'])

    for x in X_train:
        print(x.shape)

    for x in X_valid:
        print(x.shape)

    X_train = np.concatenate(X_train, axis=1)
    X_valid = np.concatenate(X_valid, axis=1)

    # check
    print('X_train shape:', X_train.shape)
    print('X_valid shape:', X_valid.shape)
    print('y_train shape:', y_train.shape)
    print('y_valid shape:', y_valid.shape)

    X_train, y_train = shuffle(X_train, y_train)
    y_train = to_categorical(y_train)
    X_valid, y_valid = shuffle(X_valid, y_valid)
    y_valid = to_categorical(y_valid)

    print('X_train shape:', X_train.shape)
    print('X_valid shape:', X_valid.shape)
    print('y_train shape:', y_train.shape)
    print('y_valid shape:', y_valid.shape)

    inputs = Input(X_train.shape[1:])
    x = Dense(2048, activation='relu')(inputs)
    x = Dropout(DROP_RATE)(x)
    x = Dense(1024, activation='relu')(x)
    predictions = Dense(CLASS_NUM, activation='softmax')(inputs)
    model = Model(inputs, predictions)
    check_point = ModelCheckpoint(filepath=os.path.join(MODEL_PATH, 'merge-model-01.hdf5'), verbose=1, save_best_only=True)
    early_stopping = EarlyStopping(monitor='val_loss', min_delta=0.0001, patience=EARLY_STOPPING_PATIENCE, verbose=1, mode='auto')
    # 创建一个 LossHistory 实例
    history = LossHistory()

    model.compile(loss='binary_crossentropy', optimizer=Adam(lr=LEARNING_RATE), metrics=['accuracy'])
    model.fit(X_train, y_train, epochs=EPOCHS, batch_size=BATCH_SIZE, validation_data=(X_valid, y_valid), callbacks=[early_stopping, check_point, history])

    # 绘制 loss 曲线和 batch 曲线
    history.loss_plot('batch', os.path.join(RESULT_PATH, 'merge_all_loss_batch.png'))
    history.acc_plot('batch', os.path.join(RESULT_PATH, 'merge_all_acc_batch.png'))
    history.loss_plot('epoch', os.path.join(RESULT_PATH, 'merge_all_loss_epoch.png'))
    history.acc_plot('epoch', os.path.join(RESULT_PATH, 'merge_all_acc_epoch.png'))

if __name__ == '__main__':
    mergeFinetuneModel()
